/*
 * CudaAesOperationsTest.cpp
 *
 *  Created on: Feb 21, 2010
 *      Author: Łukasz Spintzyk
 */
//#define DBG_NO_CUDA


#include "CudaAESOperations.h"
#include "CudaAESInvOperations.h"
#include "aesUtils.h"

#include <iostream>
#include <iomanip>

#include <cppunit/TestRunner.h>
#include <cppunit/TestResult.h>
#include <cppunit/TestResultCollector.h>
#include <cppunit/extensions/HelperMacros.h>
#include <cppunit/BriefTestProgressListener.h>
#include <cppunit/extensions/TestFactoryRegistry.h>


#include "AesKeyTestVectors.h"

class TestAesKey: public CPPUNIT_NS::TestCase
{
CPPUNIT_TEST_SUITE(TestAesKey);
		CPPUNIT_TEST(invExpandKeyTest4keyx4state);
		CPPUNIT_TEST(expandKeyTestZERO4keyx4state);
		CPPUNIT_TEST(expandKeyTestZERO4keyx8state);
		CPPUNIT_TEST(expandKeyTestFF6keyx4state);
		CPPUNIT_TEST(expandKeyTestRand8keyx4state);
		CPPUNIT_TEST(expandKeyTestRand8keyx6state);
		CPPUNIT_TEST(expandKeyTestRand8keyx8state);
	CPPUNIT_TEST_SUITE_END();

public:
	void setUp(void)
	{
		d_keySize = 4*16;
		d_key = new byte[d_keySize];
		d_expandedKey = new byte[(8*15)*4];
		d_expandedKey2 = new byte[(8*15)*4];
		memset(d_expandedKey, 0, (8*15)/4);
		setTestBlock(d_key, 8, 4);

		aesEncryptInit();
	}
	void tearDown(void)
	{
		delete[] d_key;
		delete[] d_expandedKey;
		delete[] d_expandedKey2;
	}

protected:
	void expandKeyTestZERO4keyx4state(void)
	{
		memset(d_key, 0, d_keySize);

		expandKey(d_key, NB_4, d_expandedKey, NB_4,0);
		CPPUNIT_ASSERT(memcmp(d_expandedKey, c_expandedKeyZero, 16*11) ==0 );
	}

	void invExpandKeyTest4keyx4state(void)
	{
		byte c_key[16] =
		{
				0x0f, 0x15 , 0x71 , 0xc9
				, 0x47 , 0xd9 , 0xe8 , 0x59
				, 0x0c , 0xb7 , 0xad , 0xd6
				, 0xaf , 0x7f , 0x67 , 0x98
		};

		memset(d_key, 0, d_keySize);
		expandKey(c_key, NB_4, d_expandedKey, NB_4,0);
		memcpy(d_expandedKey2, d_expandedKey, (8*15)*4);
		invExpandedKey(d_expandedKey2, NB_4, NB_4);


		for(int i = 1; i < NR[NB_4][NB_4]; ++i)
		{
			invMixColumn(d_expandedKey+i*4*4, NB_4);
		}
		//std::cout << "\ninv Expanded key"<<std::endl;
		//printBlock(d_expandedKey, 11*4,4);
		//std::cout << "\nCuda inv Expanded key"<<std::endl;
		//		printBlock(d_expandedKey2, 11*4,4);
		CPPUNIT_ASSERT(memcmp(d_expandedKey, d_expandedKey2, 11*16) ==0 );//first block is the same


	}

	void expandKeyTestZERO4keyx8state(void)
	{
		memset(d_key, 0, d_keySize);

		expandKey(d_key, NB_4, d_expandedKey, NB_8,0);
		CPPUNIT_ASSERT(memcmp(d_expandedKey, c_expandedKeyZero, 16*11) ==0 );
	}

	void expandKeyTestFF6keyx4state(void)
	{
		expandKey((byte*)c_key192FF, NB_6, d_expandedKey, NB_4,0);
		CPPUNIT_ASSERT(memcmp(d_expandedKey, c_expandedKey192FF, 13*16) ==0 );
	}

	void expandKeyTestRand8keyx4state(void)
	{
		expandKey((byte*)c_key256, NB_8, d_expandedKey, NB_4,0);
		CPPUNIT_ASSERT(memcmp(d_expandedKey, c_expandedKey256, 15*16) ==0 );
	}

	void expandKeyTestRand8keyx6state(void)
	{
		expandKey((byte*)c_key256, NB_8, d_expandedKey, NB_6,0);
		CPPUNIT_ASSERT(memcmp(d_expandedKey, c_expandedKey256, 15*16) ==0 );
	}

	void expandKeyTestRand8keyx8state(void)
	{
		expandKey((byte*)c_key256, NB_8, d_expandedKey, NB_8,0);
		CPPUNIT_ASSERT(memcmp(d_expandedKey, c_expandedKey256, 15*16) ==0 );
	}
private:
	void setTestBlock(byte* block, int w, int h, int seed=1)
	{
		int k = seed;
		for (int i = 0; i < w; ++i)
		{
			for (int j = 0; j < h; ++j)
			{
				block[i * h + j] = k++;
			}
		}
	}

	void printBlock(byte* block, int w, int h)
	{
		for (int j = 0; j < h; ++j)
		{
			for (int i = 0; i < w; ++i)
			{
				std::cout << std::hex << std::setw(4) << (unsigned int) block[i * h + j]
						<< " ";
			}
			std::cout << std::endl;
		}
	}
	byte* d_key;
	byte* d_expandedKey;
	byte* d_expandedKey2;
	int d_keySize;
};

CPPUNIT_TEST_SUITE_REGISTRATION(TestAesKey);

int main(int ac, char **av)
{
	//--- Create the event manager and test controller
	CPPUNIT_NS::TestResult controller;
	//--- Add a listener that colllects test result
	CPPUNIT_NS::TestResultCollector result;
	controller.addListener(&result);
	//--- Add a listener that print dots as test run.
	CPPUNIT_NS::BriefTestProgressListener progress;
	controller.addListener(&progress);
	//--- Add the top suite to the test runner
	CPPUNIT_NS::TestRunner runner;
	runner.addTest(CPPUNIT_NS::TestFactoryRegistry::getRegistry().makeTest());
	runner.run(controller);

	return result.wasSuccessful() ? 0 : 1;
}
